package com.centnet.base.bigdata.solr;

/**
 * Created by veblen on 2017/4/17.
 */

import com.google.common.base.Joiner;
import org.apache.commons.lang.StringUtils;
import org.apache.solr.client.solrj.SolrClient;
import org.apache.solr.client.solrj.SolrQuery;
import org.apache.solr.client.solrj.SolrServerException;
import org.apache.solr.client.solrj.beans.DocumentObjectBinder;
import org.apache.solr.client.solrj.impl.CloudSolrClient;
import org.apache.solr.client.solrj.request.UpdateRequest;
import org.apache.solr.client.solrj.response.FacetField;
import org.apache.solr.client.solrj.response.GroupResponse;
import org.apache.solr.client.solrj.response.QueryResponse;
import org.apache.solr.client.solrj.response.UpdateResponse;
import org.apache.solr.common.SolrDocument;
import org.apache.solr.common.SolrDocumentList;
import org.apache.solr.common.SolrInputDocument;
import org.apache.solr.common.params.CommonParams;
import org.apache.solr.common.params.CursorMarkParams;
import org.apache.solr.common.params.FacetParams;
import org.apache.solr.common.params.GroupParams;
import org.apache.solr.common.util.NamedList;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.IOException;
import java.util.*;


public class SolrUtil {


    public static final Logger logger = LoggerFactory.getLogger(SolrUtil.class);
    /**
     * 特殊字符
     */
    public static final String SPECIAL_SYMBOLS = "([+\\-&|!(){}\\[\\]\\^\"~\\*\\?\\\\:]{1})";

    /**
     * solr的数据转换类
     */
    public static final DocumentObjectBinder BINDER = new DocumentObjectBinder();


    /**
     * 特殊字符处理
     *
     * @param fields 字段
     * @return
     */
    public static String replaceSpecialSymbols(String fields) {

        //special chars in solr query expression:+ - && || ! ( ) { } [ ] ^ " ~ * ? : \
        //'''P'''是一个占位符，这里用它为了避免\$，使分组捕获被转义，在下一语句中会将'''P'''替换为\
        fields = fields.replaceAll(SPECIAL_SYMBOLS, "'''P'''$1");
        fields = fields.replace("'''P'''", "\\");
        return fields;
    }

    public static String replaceSpecialSym(String field){
        String ss = "+,-,&&,||,!,(,),{,},[,],^,\",~,*,?,:,\\, ";
        final String[] target = {field};
        List<String> ssList = Arrays.asList(ss.split(","));
        ssList.stream().forEach(s -> {
            target[0] = target[0].replace(s,"\\"+s);
        });
        return target[0];
    }


    private static final int timeout = 1000 * 60 * 15;

    /**
     * key是solr的index名称，value是solr客户端
     */
    private static final Map<String, SolrClient> SOLR_CLIENT_MAP = new HashMap<>();

    /**
     * 获取solr客户端
     *
     * @param collectionName 索引的名称，如果是集群模式部署的solr，这里指的collections
     * @return
     */
    public static SolrClient getSolr(String collectionName) {
        SolrClient solr = SOLR_CLIENT_MAP.get(collectionName);
        if (solr == null) {
            CloudSolrClient solrClient = new CloudSolrClient.Builder().withZkHost(SolrConfig.SOLR_HOST).build();
            solrClient.setDefaultCollection(collectionName);
            solrClient.setZkClientTimeout(timeout);
            solrClient.setZkConnectTimeout(timeout);
            solrClient.setSoTimeout(timeout);
            SOLR_CLIENT_MAP.put(collectionName, solrClient);
            return solrClient;
        }
        return solr;
    }

    /**
     * 添加单条记录
     *
     * @param map
     * @throws SolrServerException
     * @throws IOException
     */
    public static void add(String collectionName, Map<String, Object> map) throws IOException, SolrServerException {

        if (null != map && !map.isEmpty()) {
            SolrInputDocument doc = new SolrInputDocument();
            Set<String> set = map.keySet();
            for (String key : set) {
                Object value = map.get(key);
                doc.addField(key, value);
            }
            SolrClient solr = getSolr(collectionName);
            solr.add(doc);
        }
    }

    /**
     * 批量添加
     *
     * @param list
     * @param collectionName
     */
    public static void addList(String collectionName, List<Map<String, Object>> list) throws IOException, SolrServerException {
        if (null != list && list.size() > 0) {
            List<SolrInputDocument> docList = new ArrayList<SolrInputDocument>();

            for (Map<String, Object> map : list) {
                SolrInputDocument doc = new SolrInputDocument();
                Set<String> set = map.keySet();
                for (String key : set) {
                    Object value = map.get(key);
                    doc.addField(key, value);
                }
                docList.add(doc);
            }

            SolrClient solr = getSolr(collectionName);
            solr.add(docList);
        }
    }

    /**
     * 单条添加
     *
     * @param collectionName
     * @param message
     */
    public static void addBean(String collectionName, Object message) throws IOException, SolrServerException {
        if (null != message) {
            SolrClient solr = getSolr(collectionName);
            solr.addBean(message);
        }
    }

    /**
     * 批量添加
     *
     * @param collectionName
     * @param list
     */
    public static void addBeans(String collectionName, List list) throws IOException, SolrServerException {
        if (null != list && list.size() > 0) {
            SolrClient solr = getSolr(collectionName);
            solr.addBeans(list);
        }
    }

    /**
     * 添加索引
     *
     * @param collection
     * @param docs
     */
    public static void add(String collection, Collection<SolrInputDocument> docs) throws IOException, SolrServerException {
        if (null != docs && docs.size() > 0) {
            SolrClient solr = getSolr(collection);
            solr.add(collection, docs);
        }
    }

    /**
     * 更新列,
     *
     * @param collection
     * @param docs
     * @throws IOException
     * @throws SolrServerException
     */
    public static void update(String collection, List<SolrInputDocument> docs) throws IOException, SolrServerException {
        // TODO: 2017/10/23 不会报错，但是没有生效，此处代码暂时保留。
        UpdateRequest updateRequest = new UpdateRequest();
        updateRequest.setAction(UpdateRequest.ACTION.COMMIT, false, false);
        updateRequest.add(docs);
        UpdateResponse process = updateRequest.process(getSolr(collection));
        int status = process.getStatus();
        NamedList<Object> response = process.getResponse();
        for (Map.Entry<String, Object> stringObjectEntry : response) {
            System.out.println(stringObjectEntry);
        }
        System.out.println(status);
    }

    public static List<FacetField> facetSearch(String collectionName, int pageNo, int pageSize, int facetMinCount,
                                               List<String> queryStrs,
                                               String[] facetFields,
                                               String sort) throws IOException, SolrServerException {
        SolrClient solr = getSolr(collectionName);

        SolrQuery query = new SolrQuery();

        setQuery(query, queryStrs);

        // 设置通过facet查询为true，表示查询时使用facet机制
        query.setFacet(true);

        // 设置facet的字段名称是
        query.set(FacetParams.FACET_FIELD, facetFields);
        if (pageNo != 0 && pageSize != 0) {
            query.setFacetLimit(pageSize);
            query.set(FacetParams.FACET_OFFSET, (pageNo - 1) * pageSize);
        }
        // 设置统计后数量大于等于1才返回

        query.setFacetMinCount(facetMinCount);
//            query.setFacetMissing(true);
        // 设置返回doc文档数据，因只需要数量，故设置为0
        query.setRows(0);
        if (StringUtils.isNotBlank(sort)) {
            query.setFacetSort(sort);
        }
        QueryResponse response = solr.query(query);
        return response.getFacetFields();
    }

    public static List<FacetField> facetSearch(String collectionName, int pageNo, int pageSize, int facetMinCount,
                                               List<String> queryStrs,
                                               String[] facetFields,
                                               String sort,String filterFields) throws IOException, SolrServerException {
        SolrClient solr = getSolr(collectionName);

        SolrQuery query = new SolrQuery();

        setQuery(query, queryStrs);

        query.setParam("fl",filterFields);
        // 设置通过facet查询为true，表示查询时使用facet机制
        query.setFacet(true);

        // 设置facet的字段名称是
        query.set(FacetParams.FACET_FIELD, facetFields);
        if (pageNo != 0 && pageSize != 0) {
            query.setFacetLimit(pageSize);
            query.set(FacetParams.FACET_OFFSET, (pageNo - 1) * pageSize);
        }
        // 设置统计后数量大于等于1才返回

        query.setFacetMinCount(facetMinCount);
//            query.setFacetMissing(true);
        // 设置返回doc文档数据，因只需要数量，故设置为0
        query.setRows(0);
        if (StringUtils.isNotBlank(sort)) {
            query.setFacetSort(sort);
        }
        QueryResponse response = solr.query(query);
        return response.getFacetFields();
    }

    /**
     * 设置查询条件，并将特殊字符转义
     *
     * @param query
     * @param queryStrs
     */
    private static void setQuery(SolrQuery query, List<String> queryStrs) {
        if (queryStrs != null && !queryStrs.isEmpty()) {
            String join = Joiner.on(" AND ").skipNulls().join(queryStrs);
            query.setQuery(join);
        } else {
            query.setQuery("*:*");
        }
    }

    /**
     * 分组
     *
     * @param collectionName 索引名称
     * @param pageNo         分页号
     * @param pageSize       分页大小
     * @param queryStrs      查询条件
     * @param groupField     分组字段
     * @param sort           排序字段
     * @return
     * @throws IOException
     * @throws SolrServerException
     */
    public static GroupResponse grouping(String collectionName, int pageNo, int pageSize,
                                         List<String> queryStrs,
                                         String[] groupField,
                                         String sort,Integer rows,Integer start,String filterFields) throws IOException, SolrServerException {
        SolrClient solr = getSolr(collectionName);

        SolrQuery query = new SolrQuery();

        setQuery(query, queryStrs);
        //是否开启分组查询
        query.set("group", true);
        //分组字段
        query.set("group.field", groupField);
        if (pageNo != 0 && pageSize != 0) {
            query.set("group.limit", pageSize);
            query.set(GroupParams.GROUP_OFFSET, (pageNo - 1) * pageSize);
        }
        if (StringUtils.isNotBlank(sort)) {
            //排序
            query.set("group.sort", sort);
        }
        query.setRows(rows);
        query.setStart(start);
        //	（true/false）等于true,将只返回最后一个group.query查询的分组数据，
        // 一定是最后一个group.query，
        // 如果最后的条件是group.field，也没有效果，同时其他分组统计条件都会失去作用
        query.set("group.main", false);
        query.set("group.format", "grouped");
        query.set("group.ngroups", "true");
        query.set("group.truncate", "true"); //default is false;

        if (filterFields != null) {
            query.setParam("fl", filterFields);
        }

        QueryResponse queryResponse = solr.query(query);
        return queryResponse.getGroupResponse();
    }

    /**
     * @param collectionName 索引名
     * @param pageNo         当期页
     * @param pageSize       分页大小
     * @param queryStrs      查询条件
     * @param sortMap        排序方式
     * @param cla            类的class
     * @param filterFields   过滤的字段
     * @param useCache       true表示使用缓存
     * @return
     */
    public static Page queryPage(String collectionName, int pageNo, int pageSize, List<String> queryStrs, Map<String, SolrQuery.ORDER> sortMap, Class cla, String[] filterFields, boolean useCache) throws IOException, SolrServerException {

        QueryResponse rsp = getPageQueryResponse(collectionName, pageNo, pageSize, queryStrs, sortMap, filterFields, useCache);
        List beans = rsp.getBeans(cla);
        long totalCount = rsp.getResults().getNumFound();
        Page page = new Page();
        page.setList(beans);
        page.setTotalCount(totalCount);
        page.setTotalPage(getTotalPage(totalCount, pageSize));
        return page;
    }

    /**
     * @param collectionName 索引名
     * @param pageNo         当期页
     * @param pageSize       分页大小
     * @param queryStrs      查询条件
     * @param sortMap        排序方式
     * @param cla            类的class
     * @param filterFields   过滤的字段
     * @param useCache       true表示使用缓存
     * @return
     */
    public static Page queryPage(String collectionName, int pageNo, int pageSize, List<String> queryStrs, Map<String, SolrQuery.ORDER> sortMap, Class cla, String[] filterFields, boolean useCache, CompletionFieldCallback callback) throws IOException, SolrServerException {

        QueryResponse rsp = getPageQueryResponse(collectionName, pageNo, pageSize, queryStrs, sortMap, filterFields, useCache);

        SolrDocumentList results = rsp.getResults();
        List beans = new ArrayList<>(results.size());
        invokeCompletionCallback(cla, callback, beans, results);
        long totalCount = results.getNumFound();
        Page page = new Page();
        page.setList(beans);
        page.setPageSize(results.size());
        page.setTotalCount(totalCount);
        page.setTotalPage(getTotalPage(totalCount, pageSize));
        return page;
    }


    /**
     * 获取分页查询的response
     *
     * @param collectionName
     * @param pageNo
     * @param pageSize
     * @param queryStrs
     * @param sortMap
     * @param filterFields
     * @param useCache
     * @return
     * @throws SolrServerException
     * @throws IOException
     */
    private static QueryResponse getPageQueryResponse(String collectionName, int pageNo, int pageSize, List<String> queryStrs, Map<String, SolrQuery.ORDER> sortMap, String[] filterFields, boolean useCache) throws SolrServerException, IOException {
        SolrQuery query = new SolrQuery();
        query.setStart((pageNo - 1) * pageSize);
        query.setRows(pageSize);
        setQuery(query, queryStrs);

        if (sortMap != null && !sortMap.isEmpty()) {
            sortMap.forEach((k, v) -> {
                query.addSort(k, v);
            });
        }
        if (filterFields != null && filterFields.length > 0) {
            query.setParam("fl", filterFields);
        }
        query.set(CommonParams.CACHE, String.valueOf(useCache));
        SolrClient solr = getSolr(collectionName);
        return solr.query(query);
    }


    /**
     * 对index进行count查询
     *
     * @param collectionName 索引名称
     * @param queryStrs      solr查询条件
     * @param uniqueField    唯一字段的名称，通常指主键的名称
     * @return
     */
    public static long count(String collectionName, List<String> queryStrs, String uniqueField) throws IOException, SolrServerException {
        SolrQuery query = new SolrQuery();
        setQuery(query, queryStrs);
        query.setRows(0);
        if (StringUtils.isNotBlank(uniqueField)) {
            query.setParam("fl", uniqueField);
        }
        SolrClient solr = getSolr(collectionName);
        QueryResponse rsp = solr.query(query);
        return rsp.getResults().getNumFound();
    }

    /**
     * 游标查询
     *
     * @param <T>
     * @param collectionName  索引名称
     * @param queryStrs       查询条件
     * @param pageSize        分页大小
     * @param cursorRowsLimit 行数限制，超过后立即返回已经查询到的数据
     * @param cla             索引对应的类
     * @param sortClause      排序的方式及字段
     * @return
     */
    public static <T> List<T> cursorQuery(String collectionName, List<String> queryStrs, int pageSize, int cursorRowsLimit, Class<T> cla, SolrQuery.SortClause sortClause, CompletionFieldCallback completionFieldCallback) throws IOException, SolrServerException {
        SolrQuery query = new SolrQuery();
        setQuery(query, queryStrs);
        query.setRows(pageSize);
        query.setSort(sortClause);
        String cursorMark = CursorMarkParams.CURSOR_MARK_START;
        boolean done = false;
        SolrClient solr = getSolr(collectionName);
        List<T> list = new LinkedList<>();
        int resultSize = 0;
        while (!done && resultSize < cursorRowsLimit) {

            //计算最后一次需要查询的量
            if (cursorRowsLimit - list.size() < pageSize) {
                query.setRows(cursorRowsLimit - list.size());

                //已经超过查询行的限制，不用再进行查询了。
                done = true;
            }

            //执行solr的查询
            query.set(CursorMarkParams.CURSOR_MARK_PARAM, cursorMark);
            QueryResponse rsp = solr.query(query);


            SolrDocumentList results = rsp.getResults();
            resultSize += results.size();

            invokeCompletionCallback(cla, completionFieldCallback, list, results);

            String nextCursorMark = rsp.getNextCursorMark();
            if (cursorMark.equals(nextCursorMark)) {
                done = true;
            }
            cursorMark = nextCursorMark;
        }
        return list;
    }

    /**
     * 执行补全字段的回调函数
     *
     * @param cla
     * @param completionFieldCallback
     * @param list
     * @param results
     * @param <T>
     * @return
     */
    private static <T> boolean invokeCompletionCallback(Class<T> cla, CompletionFieldCallback completionFieldCallback, List<T> list, SolrDocumentList results) {
        boolean completionAll = false;
        if (completionFieldCallback != null && !results.isEmpty()) {
            completionAll = completionFieldCallback.completionAll(results);
        }

        if (!completionAll) {
            //通过回调函数补全字段，这样减少在查询完成后再次遍历结果
            for (SolrDocument result : results) {
                T bean = BINDER.getBean(cla, result);
                if (completionFieldCallback != null) {
                    completionFieldCallback.completion(bean);
                }
                list.add(bean);
            }
        }
        return completionAll;
    }


    public static class Page<T> {
        private int pageNo;
        private int pageSize;
        private List<T> list;
        private long totalPage;
        private long totalCount;

        public int getPageNo() {
            return pageNo;
        }

        public void setPageNo(int pageNo) {
            this.pageNo = pageNo;
        }

        public int getPageSize() {
            return pageSize;
        }

        public void setPageSize(int pageSize) {
            this.pageSize = pageSize;
        }

        public List<T> getList() {
            return list;
        }

        public void setList(List<T> list) {
            this.list = list;
        }

        public long getTotalPage() {
            return totalPage;
        }

        public void setTotalPage(long totalPage) {
            this.totalPage = totalPage;
        }

        public long getTotalCount() {
            return totalCount;
        }

        public void setTotalCount(long totalCount) {
            this.totalCount = totalCount;
        }
    }

    /**
     * 查找list
     *
     * @param collectionName 索引名称
     * @param sortMap        排序字段
     * @param cla            实体类
     * @param queryStrs      查询条件
     * @param rows           行数限制
     * @return
     */
    public static <T> List<T> query(String collectionName, Map<String, SolrQuery.ORDER> sortMap, Class<T> cla, List<String> queryStrs, Integer rows) throws IOException, SolrServerException {
        SolrQuery query = new SolrQuery();
        setQuery(query, queryStrs);

        if (sortMap != null && !sortMap.isEmpty()) {
            sortMap.forEach((k, v) -> {
                query.addSort(k, v);
            });
        }
        if (rows != null) {
            query.setRows(rows);
        }
        SolrClient solr = getSolr(collectionName);
        QueryResponse rsp = solr.query(query);
        List beans = rsp.getBeans(cla);
        return beans;
    }

    /**
     * 查找list
     *
     * @param collectionName 索引名称
     * @param sortMap        排序字段
     * @param cla            实体类
     * @param queryStrs      查询条件
     * @param <T>
     * @return
     */
    public static <T> List<T> query(String collectionName, Map<String, SolrQuery.ORDER> sortMap, Class<T> cla, List<String> queryStrs) throws IOException, SolrServerException {
        return query(collectionName, sortMap, cla, queryStrs, null);
    }


    /**
     * 计算总页数
     *
     * @param totalCount
     * @param pageSize
     * @return
     */
    private static long getTotalPage(long totalCount, int pageSize) {
        long totalPage = totalCount / pageSize;
        if (totalCount % pageSize != 0 || totalPage == 0) {
            totalPage++;
        }
        if (totalCount == 0) {
            totalPage = 0;
        }
        return totalPage;
    }


    /**
     * 分组查询
     *
     * @param collectionName 索引名称
     * @param fields
     * @param facetFields
     * @param keyword
     * @return
     */
    public static void queryGroup(String collectionName, String[] fields, String[] facetFields, String keyword) throws IOException, SolrServerException {
        StringBuffer queryStr = new StringBuffer();

        if (fields != null && fields.length > 0) {
            for (String field : fields) {
                queryStr.append(field).append(" ");
            }
        } else {
            queryStr.append("*");
        }

        SolrQuery query = new SolrQuery(keyword);

        query.setParam("qf", queryStr.toString().trim())
                .setParam("defType", "edismax") //qf方式查询，需要用edismax解析器
                .setIncludeScore(false)    //是否按每组数量高低排序
                .setFacet(true)            //启用分组
                .setRows(0)                //设置返回结果条数，如果你时分组查询，你就设置为0
                .setFacetMinCount(1)    //只显示大于等于1的记录
                .setFacetLimit(5)        //限制每次返回结果数
                .addFilterQuery("proPrice:[500 TO 1000]");

        for (String facetField : facetFields) {
            query.addFacetField(facetField.trim());//对哪个字段做分组统计
        }

        QueryResponse resp = getSolr(collectionName).query(query);

        List<FacetField> facets = resp.getFacetFields();
        for (FacetField ff : facets) {
            System.out.println(ff.getName() + ":" + ff.getValueCount());

            List<FacetField.Count> list = ff.getValues();
            for (FacetField.Count c : list) {
                System.out.println(c.getName() + ":" + c.getCount());
            }
        }
    }

    /**
     * 硬提交，写到硬盘
     * 在代码中，一般不显示调用，可以在配置文件中配置每隔多少时间自动调用一次
     *
     * @param collectionName
     */
    public static void commit(String collectionName) throws IOException, SolrServerException {
        getSolr(collectionName).commit();
    }

    /**
     * 软提交，写到内存
     * 在代码中，一般不显示调用，可以在配置文件中配置每隔多少时间自动调用一次
     *
     * @param collectionName
     */
    public static void softCommit(String collectionName) throws IOException, SolrServerException {
        getSolr(collectionName).commit(true, true, true);
    }

    /**
     * 判断rowkey是否存在
     *
     * @param rowKey
     * @return
     */
    public static boolean isRowKeyExists(String collectionName, String rowKey) {
        SolrQuery query = new SolrQuery();
        query.setQuery("rowKey:" + rowKey);
        query.setParam("fl", "rowKey");
        QueryResponse rsp = null;
        try {
            rsp = getSolr(collectionName).query(query);
            return rsp.getResults().size() > 0;
        } catch (SolrServerException e) {
            throw new RuntimeException(e);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    /**
     * 索引段优化
     * 将磁盘上的多个索引段合成一个大的索引段
     *
     * @param collectionName
     */
    public static void optimize(String collectionName) throws IOException, SolrServerException {
        getSolr(collectionName).optimize();
    }


    /**
     * 生成按区间查询的字符串
     *
     * @param field    字段名
     * @param startVal 开始值
     * @param endVal   结束值
     * @param isEq     true表示只有开始或结束值的时候，做相等运算还是大于小于。
     * @return
     */
    public static String genIntervalQueryStr(String field, String startVal, String endVal, boolean isEq) {
        if ("null".equals(startVal)) {
            startVal = null;
        }
        if ("null".equals(endVal)) {
            endVal = null;
        }

        String str = null;
        if (StringUtils.isNotBlank(startVal) && StringUtils.isNotBlank(endVal)) {
            str = field + ":[ " + startVal + " TO " + endVal + "]";
        } else if (StringUtils.isNotBlank(startVal) && StringUtils.isBlank(endVal)) {
            if (isEq) {
                str = field + ":" + startVal;
            } else {
                str = field + ":[" + startVal + " TO *]";
            }
        } else if (StringUtils.isNotBlank(endVal) && StringUtils.isBlank(startVal)) {
            if (isEq) {
                str = field + ":" + endVal;
            } else {
                str = field + ":[* TO " + endVal + "]";
            }
        }
        return str;
    }


}